import heapq
import math

class AStar:
    def __init__(self, size_cx, size_cy, grid, start, end):
        self.size_cx = size_cx
        self.size_cy = size_cy
        self.grid = grid
        self.start = start
        self.end = end
        self.path = []
        self.visited = set()
        self.open_set = []
        self.__came_from = {}
        self.__g_score = {}
        self.__f_score = {}
        self.running = False
        self.finished = False
        # 初始化网格和分数
        self.reset_grid()

    def reset_grid(self):
        # 重置算法状态
        self.path = []
        self.visited = set()  # 已探索节点集合
        self.open_set = []  # 开放集：它是一个优先队列，用来存储待探索节点。它的元素是Tuple(f值, Tuple(x坐标, y坐标))。
        self.__came_from = {}  # 记录节点的父节点（路径回溯）
        self.__g_score = {}  # 从起点到各节点的实际代价
        self.__f_score = {}  # 估计的总代价 (g_score + heuristic)，即g_score加上启发式函数估计的剩余代价。
        self.running = False
        self.finished = False

        # 初始化A*算法数据结构
        # g_score和f_score在大地图时内存可能会成为问题，例如100*100需要的是10,000*4约40KB，1000*1000需要约8MB，10,000*10,000需要约
        # 800MB。这里应适当进行优化，在性能和内存占用之间做一些平衡。
        for y in range(self.size_cy):
            for x in range(self.size_cx):
                self.__g_score[(x, y)] = float('inf')  # 初始化所有节点的g_score和f_score为无穷大
                self.__f_score[(x, y)] = float('inf')

        # 设置起点的初始值
        self.__g_score[self.start] = 0  # 起点到起点本身的实际代价是0。
        self.__f_score[self.start] = self.__heuristic(self.start, self.end) # 起点到终点估计的代价。
        heapq.heappush(self.open_set, (self.__f_score[self.start], self.start)) # 把起点加入到待探索节点中。

    @staticmethod
    def __heuristic(a, b):
        """
        启发式函数:用来估算a~b的代价。
        在A*算法中起到至关重要的作用，如果选对了事半功倍。在open set中选择下一个节点时，根据(起点到节点的当前最短距离 + a~b的估算距离)
        来取，取值最小的一个。
        在100*100小地图的情况下：使用平方差((x1-x0)^2 + (y1-y0)^2)优于使用曼哈顿距离和欧几里得距离，平方差放大启发式函数的作用（即
        在open set中选择下一个节点时减少了"起点到节点的当前最短距离"的作用）。
        在面对不同大小的地图时可能需要做对应调整。
        注意！需要针对性更改这个函数，当网格的通行代价很大时这个启发函数起到的作用非常少。
        """
        # 曼哈顿距离
        # return abs(a[0] - b[0]) + abs(a[1] - b[1])
        # 欧几里得距离
        #return math.sqrt((a[0] - b[0])**2 + (a[1] - b[1])**2)
        # 平方差：放大启发式函数的作用。
        return (a[0] - b[0])**2 + (a[1] - b[1])**2

    @staticmethod
    def __weight(weight):
        """
        当前的设定是weight值在1~127之间（方便绘图）。这里返回输入参数的平方，放大weight的影响。
        对应的启发式函数用的是差的平方，所以估算距离时“距离”的值会很大，因此这里也做一个平方，避免在计算距离时候这个分量被忽略了。
        在面对不同大小的地图时可能需要做对应调整。
        """
        return weight ** 2

    def __get_neighbors(self, pos):
        """
        返回pos位置的邻居坐标列表，如果邻居不可达则它不会出现在返回列表中。可以选择是4个还是8个方向移动，如果是八方则最多返回8个邻居。
        """
        x, y = pos
        neighbors = []

        # 四方向移动
        #directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
        # 八方向移动
        directions = [(0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1), (-1, 0), (-1, 1)]

        for dx, dy in directions:
            nx, ny = x + dx, y + dy
            if (0 <= nx < self.size_cx and 0 <= ny < self.size_cy
                    and self.grid[nx][ny][0] == 1 and self.grid[nx][ny][1] != 255):  # 通行代价不能是255
                neighbors.append((nx, ny))

        return neighbors

    def step(self):
        if not self.running or self.finished:
            return

        if not self.open_set:
            self.finished = True
            self.running = False
            return

        # 在待探索节点列表中取出f值最小的一个节点。这个有序优先堆队列中存的是Tuple(f值, Tuple(x坐标, y坐标))
        # f值：估计的总代价 (g_score + heuristic)，即从起点到各节点的实际代价加上启发式函数估计的剩余代价。
        nc = heapq.heappop(self.open_set)
        current = nc[1]
        #print("待探索节点列表中取出f值最小的一个节点:",current, nc[0])
        if current == self.end:
            # 如果触及终点，那么反向重建路径
            self.path = []
            while current in self.__came_from:  # came_from记录节点的上一个节点。
                self.path.append(current)
                current = self.__came_from[current]
            self.path.append(self.start)
            self.path.reverse()
            self.finished = True
            self.running = False
            print("-----------summery-------------")
            print("path:", len(self.path))
            print("visited:", len(self.visited))
            print("open_set:", len(self.open_set))
            print("__came_from:", len(self.__came_from))
            print("__g_score:", len(self.__g_score))
            print("__f_score:", len(self.__f_score))
            return

        # 将f值最小的节点从待探索节点列表中取出以后，探索这个节点的所有邻居，这个节点将被标记为已探索。
        self.visited.add(current)
        for neighbor in self.__get_neighbors(current):
            if neighbor in self.visited:  # 如果这个节点已探索，跳过。
                continue
            # g_score是从起点到各节点的实际代价
            x, y = neighbor[0], neighbor[1]
            # 测试：遇到邻居的代价不为1时中断
            # if self.grid[x][y][1] != 1:
            #     print("当前", current, "邻居({0},{1})代价:{2}".format(x, y, self.grid[x][y][1] ** 2))
            tentative_g_score = self.__g_score[current] + self.__weight(self.grid[x][y][1]) # 当前点到它的每个邻居的代价。
            # 在程序开始时，会初始化所有节点到起点的代价g_score和f_score(经过此点到终点的预估代价)为无穷大（float('inf')）。
            # 所以如果某个点未曾访问过，它的g_score总是大于tentative_g_score。而如果该点曾经访问过但是前一个路径的代价大于当前路径的代价，
            # 那么用新的父节点代替旧值，并更新g_score和f_score。
            if tentative_g_score < self.__g_score[neighbor]:
                self.__came_from[neighbor] = current
                self.__g_score[neighbor] = tentative_g_score
                self.__f_score[neighbor] = tentative_g_score + self.__heuristic(neighbor, self.end)
                # 如果邻居节点在开放集中，检查开放集中它的f_score，如果当前f_score小于他则更新；
                # 如果邻居不在开放集中，则添加它。
                found = False
                for i, (f, pos) in enumerate(self.open_set):
                    if pos == neighbor:
                        found = True
                        if f > self.__f_score[neighbor]:
                            self.open_set[i] = (self.__f_score[neighbor], neighbor)
                            # print("open set: update ", neighbor, self.__f_score[neighbor])
                            heapq.heapify(self.open_set)
                        break
                if not found:
                    heapq.heappush(self.open_set, (self.__f_score[neighbor], neighbor))
                    # print("open set: push ", neighbor, self.__f_score[neighbor])
